import unittest
from datetime import datetime

import pandas as pd

from controller.ai.report import calculate_equipment_valid_days


class TestCalculateEquipmentValidDays(unittest.TestCase):

    def setUp(self):
        # 设置测试用的统计时间范围
        self.start_time = datetime(2023, 1, 1)
        self.end_time = datetime(2023, 1, 10)
        self.total_period_days = (self.end_time - self.start_time).days + 1  # 10天

    def test_device_created_before_start(self):
        # 设备在统计开始时间之前创建
        df_equipment = pd.DataFrame({
            'equipment_id': ['E001'],
            'createDate': ['2022-12-25T10:00:00Z']
        })
        result_df = calculate_equipment_valid_days(df_equipment, self.start_time, self.end_time)
        self.assertEqual(result_df.loc[0, '有效出勤天数'], self.total_period_days)  # 应该为10天

    def test_device_created_on_start(self):
        # 设备在统计开始时间当天创建
        df_equipment = pd.DataFrame({
            'equipment_id': ['E002'],
            'createDate': ['2023-01-01T00:00:00Z']
        })
        result_df = calculate_equipment_valid_days(df_equipment, self.start_time, self.end_time)
        self.assertEqual(result_df.loc[0, '有效出勤天数'], self.total_period_days)  # 应该为10天

    def test_device_created_in_middle(self):
        # 设备在统计周期中间创建
        df_equipment = pd.DataFrame({
            'equipment_id': ['E003'],
            'createDate': ['2023-01-05T12:00:00Z']
        })
        result_df = calculate_equipment_valid_days(df_equipment, self.start_time, self.end_time)
        # 从1月5日到1月10日，共6天 (5, 6, 7, 8, 9, 10)
        self.assertEqual(result_df.loc[0, '有效出勤天数'], 6)

    def test_device_created_on_end(self):
        # 设备在统计结束时间当天创建
        df_equipment = pd.DataFrame({
            'equipment_id': ['E004'],
            'createDate': ['2023-01-10T08:00:00Z']
        })
        result_df = calculate_equipment_valid_days(df_equipment, self.start_time, self.end_time)
        self.assertEqual(result_df.loc[0, '有效出勤天数'], 1)  # 1月10日当天，1天

    def test_device_created_after_end(self):
        # 设备在统计结束时间之后创建
        df_equipment = pd.DataFrame({
            'equipment_id': ['E005'],
            'createDate': ['2023-01-11T15:00:00Z']
        })
        result_df = calculate_equipment_valid_days(df_equipment, self.start_time, self.end_time)
        self.assertEqual(result_df.loc[0, '有效出勤天数'], 0)

    def test_device_with_invalid_create_date(self):
        # 设备创建日期无效
        df_equipment = pd.DataFrame({
            'equipment_id': ['E006'],
            'createDate': ['invalid-date']
        })
        result_df = calculate_equipment_valid_days(df_equipment, self.start_time, self.end_time)
        self.assertEqual(result_df.loc[0, '有效出勤天数'], self.total_period_days)  # 应该按全周期计算

    def test_multiple_devices(self):
        # 测试多个设备混合情况
        df_equipment = pd.DataFrame({
            'equipment_id': ['E007', 'E008', 'E009', 'E010'],
            'createDate': ['2022-12-01', '2023-01-03', '2023-01-12', None]
        })
        result_df = calculate_equipment_valid_days(df_equipment, self.start_time, self.end_time)

        # E007: 2022-12-01 < 2023-01-01, 全周期10天
        self.assertEqual(result_df.loc[0, '有效出勤天数'], 10)
        # E008: 2023-01-03, (10 - 3) + 1 = 8天
        self.assertEqual(result_df.loc[1, '有效出勤天数'], 8)
        # E009: 2023-01-12 > 2023-01-10, 0天
        self.assertEqual(result_df.loc[2, '有效出勤天数'], 0)
        # E010: None, 视为无效，全周期10天
        self.assertEqual(result_df.loc[3, '有效出勤天数'], 10)


if __name__ == '__main__':
    unittest.main(argv=['first-arg-is-ignored'], exit=False)